import torch
import torch.nn as nn
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler
from LSTM_one.LSTM.dataset_create import create_dataset
from LSTM_one.LSTM.lstm import LSTM

plt.rcParams['font.sans-serif']=['SimHei']  #图中字体改为黑体以兼容中文
plt.rcParams['axes.unicode_minus']=False #负号显示的问题
# 读取数据
data = pd.read_csv('../data/2206_2404新能源汽车总体销量数据_全国.csv',encoding='latin1')

# 归一化处理
scaler = MinMaxScaler(feature_range=(0, 1))
scaled_data = scaler.fit_transform(data['sale_count'].values.reshape(-1,1))

# 定义模型参数
input_size = 1
hidden_size = 120
num_layers = 3
output_size = 1
learning_rate = 0.001
num_epochs = 800

# 定义 LSTM_one 模型
lstm = LSTM(input_size, hidden_size, num_layers, output_size)

look_back = 6
# 加载保存的模型参数
lstm.load_state_dict(torch.load('../model/lstm_model.pth'))

# 将最后的12期数据作为初始输入
test_input = scaled_data[-look_back:]
predictions = []
# 预测未来六期销量，并更新下一期预测数据的输入
for i in range(6):
    # 将数据转换为张量并添加两个新维度（batch 和 feature）
    test_input = torch.from_numpy(test_input).type(torch.Tensor)
    test_input = test_input.unsqueeze(0)  # 在第一个维度和最后一个维度增加维度1
    # 模型预测
    lstm.eval()  # 设置模型为评估模式
    with torch.no_grad():
        prediction = lstm(test_input)  # 注意：使用新的维度
    prediction = prediction.numpy()[0][0]  # 转换成NumPy数组并取第一个值
    predictions.append(prediction)
    # 更新下一期预测数据的输入
    test_input = np.concatenate([test_input.squeeze().detach().numpy()[1:], [prediction]]).reshape(-1,1)  # 去掉新增维度，并在末尾添加预测值
# 反归一化，得到最终的预测结果
predictions = np.array(predictions).reshape(-1, 1)
predictions = scaler.inverse_transform(predictions)
predictions=predictions.astype(int)
# # 绘制预测结果图
plt.plot(pd.date_range(start='2024-05', end='2024-11', freq='M'),predictions,label='未来值')
plt.xlabel('时间')
plt.ylabel('销量')
plt.title('LSTM模型实际值与预测值比较')
plt.legend()
plt.show()

print('未来六期的销量预测值：\n', predictions)